import pandas as pd
import numpy as np
# import matplotlib.pyplot as plt
import util_for_output_zh
from util_set_zh_matplot import plt
import seaborn as sns
from pathlib import Path
import pdb
import os

import util_for_output_zh

# 在查看数据前设置
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)
pd.set_option('display.width', None)
pd.set_option('display.max_colwidth', None)

out_dir = 'ques1_prepare_data'
os.makedirs(out_dir , exist_ok=True)

def main():
    # stage1()
    # df_result, models = stage2()
    stage2_optim()
    # stage3()
    stage3_optim()

def stage1():
    df = pd.read_csv('ques1_prepare_data\sheet1_男胎检测数据_fillnan_sifted_outlier.csv')
    # 去除unnamed的列
    # 保存回原来的位置
    # print head
    # Remove columns that start with 'Unnamed'
    df = df.loc[:, ~df.columns.str.contains('^Unnamed')]

    # Save back to the original location
    df.to_csv('ques1_prepare_data\sheet1_男胎检测数据_fillnan_sifted_outlier.csv', index=False)

    # Print the head of the DataFrame
    print(df.head())

from sklearn.preprocessing import MinMaxScaler
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score, mean_squared_error
from sklearn.preprocessing import PolynomialFeatures
from util_for_ques0 import parse_gestational_age, calculate_last_menstrual_date

def stage2():
    df_male = pd.read_csv('ques1_prepare_data\sheet1_男胎检测数据_fillnan_sifted_outlier.csv')

    # 题1  试分析胎儿Y染色体浓度与孕妇的孕周数和BMI等指标的相关特性，给出相应的关系模型，并检验其显著性。 
    '''
    将 str 类型的 孕周数 转换为 数值 '孕周天数'
    去除 孕周天数 BMI Y染色体 三个指标
    1, 对 (孕周天数,BMI) 升序后 划分为五段 输出每一段的数据量
    2, 针对每一段内 孕周天数 BMI Y染色体  分别 归一化 获得数据
    3, 使用非线性变换，拟合每一段的结果，并输出拟合的效果 和 最终每一段的拟合方程
    '''
    # 1. 解析孕周为天数
    if '检测孕周' in df_male.columns:
        df_male['孕周天数'] = df_male['检测孕周'].apply(parse_gestational_age)
    else:
        df_male['孕周天数'] = None
    
    # 选择相关特征列
    relevant_cols = ['孕周天数', '孕妇BMI', 'Y染色体浓度']
    df_analysis = df_male[relevant_cols].copy()
    
    # 2. 按孕周天数和BMI排序，然后分成5个区段
    df_analysis = df_analysis.sort_values(by=['孕周天数', '孕妇BMI'])
    df_analysis['segment'] = pd.qcut(df_analysis.index, q=5, labels=False)
    
    # 打印每个区段的数据量
    print("每个区段的数据量:")
    print(df_analysis['segment'].value_counts().sort_index())
    print("\n")
    
    # 初始化结果存储
    segment_results = []
    
    # 3. 处理每个区段
    for seg in range(5):
        segment_df = df_analysis[df_analysis['segment'] == seg].copy()
        
        # 对特征进行归一化处理
        scaler = MinMaxScaler()
        segment_df[['孕周天数_norm', '孕妇BMI_norm']] = scaler.fit_transform(segment_df[['孕周天数', '孕妇BMI']])
        
        # 尝试多项式回归(degree=2)
        X = segment_df[['孕周天数_norm', '孕妇BMI_norm']]
        y = segment_df['Y染色体浓度']
        
        # 生成多项式特征
        poly = PolynomialFeatures(degree=2)
        X_poly = poly.fit_transform(X)
        
        # 训练线性回归模型
        model = LinearRegression()
        model.fit(X_poly, y)
        y_pred = model.predict(X_poly)
        r2 = r2_score(y, y_pred)
        
        # 存储结果
        segment_results.append({
            'segment': seg,
            'model': model,               # 回归模型
            'poly_features': poly,       # 多项式特征转换器
            'r2_score': r2,              # R平方得分
            'coefficients': model.coef_,  # 回归系数
            'intercept': model.intercept_, # 截距项
            'data_points': len(segment_df) # 数据点数量
        })
        
        # 打印区段结果
        print(f"区段 {seg+1} 结果:")
        print(f"数据点数量: {len(segment_df)}")
        print(f"R平方得分: {r2:.4f}")
        print("回归方程:")
        print(f"Y = {model.intercept_:.4f}", end="")
        for i, coef in enumerate(model.coef_):
            if i == 1:
                print(f" + {coef:.4f}*孕周天数", end="")
            elif i == 2:
                print(f" + {coef:.4f}*BMI", end="")
            elif i == 3:
                print(f" + {coef:.4f}*孕周天数²", end="")
            elif i == 4:
                print(f" + {coef:.4f}*孕周天数*BMI", end="")
            elif i == 5:
                print(f" + {coef:.4f}*BMI²", end="")
        print("\n")
    
    return df_analysis, segment_results

def stage2_optim():
    def plot_segment_results(segment_df, y_pred, model, poly, scaler, seg, save_dir='results/segment_plots'):
        """绘制单个区段的3D可视化结果"""
        # 创建保存目录
        os.makedirs(save_dir, exist_ok=True)
        
        # 创建更大的图形，并调整布局
        fig = plt.figure(figsize=(30, 15))  # 增加高度
        ax = fig.add_subplot(111, projection='3d')
        
        # 绘制真实值
        ax.scatter(segment_df['孕周天数'], segment_df['孕妇BMI'], segment_df['Y染色体浓度'],
                c='r', marker='o', label='真实值', alpha=0.7)
        
        # 绘制预测曲面
        x_surf = np.linspace(segment_df['孕周天数'].min(), segment_df['孕周天数'].max(), 20)
        y_surf = np.linspace(segment_df['孕妇BMI'].min(), segment_df['孕妇BMI'].max(), 20)
        x_surf, y_surf = np.meshgrid(x_surf, y_surf)
        
        # 准备网格数据的归一化版本
        grid_points = np.c_[x_surf.ravel(), y_surf.ravel()]
        grid_points_norm = scaler.transform(grid_points)
        
        # 生成多项式特征并预测
        grid_poly = poly.transform(grid_points_norm)
        z_pred = model.predict(grid_poly).reshape(x_surf.shape)
        
        # 绘制预测曲面
        surf = ax.plot_surface(x_surf, y_surf, z_pred, color='b', alpha=0.3)
        
        ax.set_xlabel('孕周天数', fontsize=12, labelpad=15)  # 增加标签间距
        ax.set_ylabel('孕妇BMI', fontsize=12, labelpad=15)
        ax.set_zlabel('Y染色体浓度', fontsize=12, labelpad=15)
        plt.title(f"区段 {seg+1} 真实值与预测曲面对比", fontsize=14, pad=20)  # 增加标题间距
        
        # 调整视角以获得更好的显示效果
        ax.view_init(elev=20, azim=45)
        
        # 调整子图位置和边距
        plt.subplots_adjust(left=0.05, right=0.95, bottom=0.05, top=0.9)
        
        # 保存图片时使用更大的bbox_inches参数
        plt.savefig(f'{save_dir}/segment_{seg+1}.png', dpi=300, bbox_inches='tight', pad_inches=0.5)
        plt.close()

    def print_segment_table(segment_results):
        """打印格式化的区段结果表格"""
        from tabulate import tabulate
        
        table_data = []
        for result in segment_results:  # 修正遍历方式
            seg = result['segment']
            equation = f"Y = {result['intercept']:.2f}"
            for i, coef in enumerate(result['coefficients']):
                # 跳过第一个系数，因为PolynomialFeatures会添加偏置项
                if i == 0:
                    continue
                if i == 1: 
                    equation += f" + {coef:.2f}×孕周天数"
                elif i == 2: 
                    equation += f" + {coef:.2f}×BMI"
                elif i == 3: 
                    equation += f" + {coef:.2f}×孕周天数²"
                elif i == 4: 
                    equation += f" + {coef:.2f}×孕周天数×BMI"
                elif i == 5: 
                    equation += f" + {coef:.2f}×BMI²"
            
            table_data.append([
                seg + 1,
                result['data_points'],
                f"{result['r2_score']:.4f}",
                equation
            ])
        
        print(tabulate(table_data, 
                    headers=['区段', '样本量', 'R²得分', '回归方程'],
                    tablefmt='grid',
                    stralign='center',
                    numalign='center'))

    def stage2_enhanced():
        # 1. 数据加载与预处理
        df_male = pd.read_csv('ques1_prepare_data/sheet1_男胎检测数据_fillnan_sifted_outlier.csv')
        df_male['孕周天数'] = df_male['检测孕周'].apply(parse_gestational_age)
        df_analysis = df_male[['孕周天数', '孕妇BMI', 'Y染色体浓度']].dropna()
        
        # 2. 数据分段
        df_analysis = df_analysis.sort_values(by=['孕周天数', '孕妇BMI'])
        df_analysis['segment'] = pd.qcut(df_analysis.index, q=5, labels=False)
        
        # 3. 分段建模
        segment_results = []
        for seg in range(5):
            segment_df = df_analysis[df_analysis['segment'] == seg].copy()
            
            if len(segment_df) == 0:
                print(f"警告：区段 {seg+1} 没有数据，跳过建模")
                continue
            
            # 归一化
            scaler = MinMaxScaler()
            X_norm = scaler.fit_transform(segment_df[['孕周天数', '孕妇BMI']])
            y = segment_df['Y染色体浓度']
            
            # 多项式回归
            poly = PolynomialFeatures(degree=2)
            X_poly = poly.fit_transform(X_norm)
            
            model = LinearRegression().fit(X_poly, y)
            y_pred = model.predict(X_poly)
            
            # 存储结果
            segment_results.append({
                'segment': seg,
                'model': model,
                'poly_features': poly,
                'scaler': scaler,
                'r2_score': r2_score(y, y_pred),
                'coefficients': model.coef_,
                'intercept': model.intercept_,
                'data_points': len(segment_df),
                'data': segment_df,
                'y_pred': y_pred
            })
            
            # 绘制该区段结果 - 修正参数传递
            plot_segment_results(segment_df, y_pred, model, poly, scaler, seg,out_dir)
        
        # 4. 打印汇总结果
        print("\n" + "="*60)
        print("胎儿Y染色体浓度分段回归分析报告".center(50))
        print("="*60)
        
        print(f"\n总样本量: {len(df_analysis)}")
        print("各区段数据分布:")
        print(df_analysis['segment'].value_counts().sort_index().to_string())
        
        print("\n" + "-"*60)
        print("各区段建模结果汇总:")
        print_segment_table(segment_results)
        
        # 5. 返回结果
        return {
            'data': df_analysis,
            'segment_results': segment_results,
            'summary_stats': {
                'avg_r2': np.mean([s['r2_score'] for s in segment_results]),
                'total_samples': len(df_analysis)
            }
        }
    # print( stage2_enhanced() )
    return stage2_enhanced()


def stage3():
    df_male = pd.read_csv('ques1_prepare_data\sheet1_男胎检测数据_fillnan_sifted_outlier.csv')

    # 题1  试分析胎儿Y染色体浓度与孕妇的孕周数和BMI等指标的相关特性，给出相应的关系模型，并检验其显著性。 
    '''
    将 str 类型的 孕周数 转换为 数值 '孕周天数'
    去除 孕周天数 BMI Y染色体 三个指标
    1, 对 孕周天数 BMI Y染色体  分别 归一化 获得数据
    2, 使用非线性变换，拟合结果，并输出拟合的效果 和 最终拟合方程
    '''
    # 1. 解析孕周为天数
    if '检测孕周' in df_male.columns:
        df_male['孕周天数'] = df_male['检测孕周'].apply(parse_gestational_age)
    else:
        df_male['孕周天数'] = None
    
    # 选择相关特征列
    relevant_cols = ['孕周天数', '孕妇BMI', 'Y染色体浓度']
    df_analysis = df_male[relevant_cols].copy()

    # 选择分析所需的三个关键特征
    relevant_cols = ['孕周天数', '孕妇BMI', 'Y染色体浓度']
    df_analysis = df_male[relevant_cols].dropna().copy()  # 去除缺失值
    
    # 2. 数据归一化处理（Min-Max标准化）
    scaler = MinMaxScaler()
    df_analysis[['孕周天数_norm', 'BMI_norm']] = scaler.fit_transform(df_analysis[['孕周天数', '孕妇BMI']])
    
    # 3. 准备特征和目标变量
    X = df_analysis[['孕周天数_norm', 'BMI_norm']]
    y = df_analysis['Y染色体浓度']
    
    # 4. 多项式特征扩展（非线性变换）
    poly = PolynomialFeatures(degree=2, include_bias=False)  # 二次多项式
    X_poly = poly.fit_transform(X)
    
    # 5. 建立并训练线性回归模型（使用扩展后的多项式特征）
    model = LinearRegression()
    model.fit(X_poly, y)
    
    # 6. 模型评估
    y_pred = model.predict(X_poly)
    r2 = r2_score(y, y_pred)
    mse = mean_squared_error(y, y_pred)
    
    # 7. 获取模型参数
    coef_dict = {
        'intercept': model.intercept_,
        '孕周天数': model.coef_[0],
        'BMI': model.coef_[1],
        '孕周天数²': model.coef_[2],
        '孕周天数*BMI': model.coef_[3],
        'BMI²': model.coef_[4]
    }
    
    # 8. 打印结果
    print("="*50)
    print("非线性回归分析结果（全局模型）")
    print("="*50)
    print(f"数据总量: {len(df_analysis)}")
    print(f"R²得分: {r2:.4f}")
    print(f"均方误差(MSE): {mse:.4f}")
    print("\n回归方程:")
    print(f"Y = {coef_dict['intercept']:.4f}")
    print(f" + {coef_dict['孕周天数']:.4f}×孕周天数")
    print(f" + {coef_dict['BMI']:.4f}×BMI")
    print(f" + {coef_dict['孕周天数²']:.4f}×孕周天数²")
    print(f" + {coef_dict['孕周天数*BMI']:.4f}×孕周天数×BMI")
    print(f" + {coef_dict['BMI²']:.4f}×BMI²")
    
    # 9. 返回结果
    return {
        'dataframe': df_analysis,
        'model': model,
        'poly_transformer': poly,
        'scaler': scaler,
        'metrics': {'r2': r2, 'mse': mse},
        'coefficients': coef_dict
    }

def stage3_optim():
    from mpl_toolkits.mplot3d import Axes3D
    from sklearn.preprocessing import MinMaxScaler, PolynomialFeatures
    from sklearn.linear_model import LinearRegression
    from sklearn.metrics import r2_score, mean_squared_error
    from sklearn.model_selection import cross_val_score, GridSearchCV
    from sklearn.pipeline import Pipeline
    import statsmodels.api as sm
    from util_for_ques0 import parse_gestational_age
    def safe_parse_gestational_age(ga_str):
        """安全解析孕周字符串（带异常捕获）"""
        try:
            return parse_gestational_age(ga_str)
        except (ValueError, AttributeError) as e:
            print(f"警告：解析孕周时出错 - {str(e)}")
            return np.nan

    def plot_3d_results(X, y_true, y_pred):
        """3D可视化真实值与预测值"""
        fig = plt.figure(figsize=(12, 10))
        ax = fig.add_subplot(111, projection='3d')
        
        # 绘制真实值（红色）
        ax.scatter(X.iloc[:, 0], X.iloc[:, 1], y_true, 
                c='r', marker='o', alpha=0.6, label='真实值')
        
        # 绘制预测值（蓝色）
        ax.scatter(X.iloc[:, 0], X.iloc[:, 1], y_pred,
                c='b', marker='^', alpha=0.6, label='预测值')
        
        ax.set_xlabel('孕周天数', fontsize=12)
        ax.set_ylabel('孕妇BMI', fontsize=12)
        ax.set_zlabel('Y染色体浓度', fontsize=12)
        plt.title("真实值与预测值对比", fontsize=14)
        plt.legend()
        plt.tight_layout()
        # 调整视角以便更好观察
        ax.view_init(elev=25, azim=45)  # 仰角25度，方位角45度
        
        # 自动调整布局
        plt.tight_layout()
        
        # 图片保存逻辑
        # 自动生成文件名（含时间戳）
        import time
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        filename = f"Y_chromosome_prediction_{timestamp}.png"
        full_path = f"{out_dir}/{filename}"
        
        # 保存图片（DPI=300保证印刷质量）
        plt.savefig(full_path, dpi=300, bbox_inches='tight', facecolor='white')
        print(f"可视化结果已保存至: {full_path}")

    def stage3_optimized():
        """优化后的完整分析流程"""
        # ==================== 1. 数据加载与校验 ====================
        try:
            df_male = pd.read_csv('ques1_prepare_data/sheet1_男胎检测数据_fillnan_sifted_outlier.csv')
        except FileNotFoundError:
            raise FileNotFoundError("数据文件未找到，请检查路径")
        
        # 关键字段校验
        required_cols = ['检测孕周', '孕妇BMI', 'Y染色体浓度']
        missing_cols = [col for col in required_cols if col not in df_male.columns]
        if missing_cols:
            raise ValueError(f"数据缺少必要字段: {missing_cols}")

        # ==================== 2. 数据预处理 ====================
        # 孕周解析（安全版本）
        df_male['孕周天数'] = df_male['检测孕周'].apply(safe_parse_gestational_age)
        
        # 选择特征并清除缺失值
        df_analysis = df_male[['孕周天数', '孕妇BMI', 'Y染色体浓度']].dropna()
        print(f"有效样本量: {len(df_analysis)} (已丢弃{len(df_male)-len(df_analysis)}个缺失样本)")
        
        # ==================== 3. 特征工程 ====================
        # 分离特征和目标
        X = df_analysis[['孕周天数', '孕妇BMI']]
        y = df_analysis['Y染色体浓度']
        
        # 特征归一化
        scaler = MinMaxScaler()
        X_scaled = scaler.fit_transform(X)
        
        # ==================== 4. 模型训练与调优 ====================
        # 构建管道（多项式特征+线性回归）
        pipe = Pipeline([
            ('poly', PolynomialFeatures(include_bias=False)),
            ('model', LinearRegression())
        ])
        
        # 网格搜索选择最佳多项式阶数
        param_grid = {'poly__degree': [1, 2, 3]}
        grid_search = GridSearchCV(pipe, param_grid, cv=5, scoring='r2')
        grid_search.fit(X_scaled, y)
        
        # 获取最佳模型
        best_model = grid_search.best_estimator_
        best_degree = grid_search.best_params_['poly__degree']
        print(f"\n最佳多项式阶数: {best_degree}")
        print(f"交叉验证平均R²: {grid_search.best_score_:.4f}")
        
        # ==================== 5. 模型评估 ====================
        # 生成多项式特征（使用最佳阶数）
        poly = PolynomialFeatures(degree=best_degree, include_bias=False)
        X_poly = poly.fit_transform(X_scaled)
        
        # 预测与评估
        y_pred = best_model.predict(X_scaled)
        final_r2 = r2_score(y, y_pred)
        final_mse = mean_squared_error(y, y_pred)
        
        # ==================== 6. 结果分析 ====================
        # 获取模型系数
        coef_dict = {
            'intercept': best_model.named_steps['model'].intercept_,
            **dict(zip(poly.get_feature_names_out(['孕周天数', 'BMI']),
                    best_model.named_steps['model'].coef_))
        }
        
        # 统计显著性检验
        X_sm = sm.add_constant(X_poly)
        sm_model = sm.OLS(y, X_sm).fit()
        
        # ==================== 7. 结果展示 ====================
        print("\n" + "="*50)
        print("非线性回归分析最终报告")
        print("="*50)
        print(f"最优模型R²: {final_r2:.4f}")
        print(f"均方误差(MSE): {final_mse:.4f}")
        
        print("\n回归系数及统计显著性:")
        print(sm_model.summary())
        
        print("\n回归方程:")
        equation = f"Y = {coef_dict['intercept']:.4f}"
        for name, coef in coef_dict.items():
            if name != 'intercept':
                equation += f" + {coef:.4f}*{name}"
        print(equation)
        
        # 可视化
        plot_3d_results(X, y, y_pred)
        
        # ==================== 8. 返回结果 ====================
        return {
            'data': df_analysis,
            'best_model': best_model,
            'scaler': scaler,
            'poly_transformer': poly,
            'metrics': {
                'r2': final_r2,
                'mse': final_mse,
                'best_degree': best_degree
            },
            'coefficients': coef_dict,
            'stats_model': sm_model
        }
    # print( stage3_optimized() )
    return stage3_optimized()

if __name__=='__main__':
    main()